/-
Copyright (c) 2025 Christian Krause. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Christian Krause
-/
module

public import Mathlib.GroupTheory.FreeGroup.Reduce
public import Mathlib.GroupTheory.GroupAction.Defs

/-!
For any `w : α × Bool`, `FreeGroup.startsWith w` is the set of all elemenents of `FreeGroup α` that
start with `w`.

The main theorem `Orbit.duplicate` proves that applying `w⁻¹` to the orbit of `x` under the action
of `FreeGroup.startsWith w` yields the orbit of `x` under the action of `FreeGroup.startsWith v`
for every `v ≠ w⁻¹` (and the point `x`).
-/

@[expose] public section

variable {α X : Type*} [DecidableEq α]

namespace FreeGroup

/--
All elements of the free Group that start with a certain letter.
-/
def startsWith (w : α × Bool) := {g : FreeGroup α | (FreeGroup.toWord g)[0]? = some w}

/--
The neutral element is not contained in one of the startsWith sets.
-/
theorem startsWith.ne_one {w : α × Bool} (g : FreeGroup α) (h : g ∈ FreeGroup.startsWith w) :
    g ≠ 1 := fun h1 ↦ by simp [h1, startsWith, FreeGroup.toWord_one] at h

@[simp]
lemma startsWith.disjoint_iff_ne {w w' : α × Bool} :
    Disjoint (startsWith w) (startsWith w') ↔ w ≠ w'  := by
  simp_all only [ne_eq, startsWith, Set.disjoint_iff_inter_eq_empty, Set.ext_iff, Set.mem_inter_iff,
    Set.mem_setOf_eq, Set.mem_empty_iff_false, iff_false, not_and, Option.some.injEq]
  exact Iff.intro (fun h ↦ h (mk [w]) (by simp)) (by grind)

lemma startsWith.Injective : @startsWith α _|>.Injective := fun a b h ↦ by
  simp only [startsWith, Set.ext_iff, Set.mem_setOf_eq] at h
  simpa using h (mk [a])

theorem startsWith_mk_mul {w : α × Bool} (g : FreeGroup α)
    (h : ¬ g ∈ startsWith (w.1, !w.2)) : mk [w] * g ∈ startsWith w := by
  by_cases hC : 0 < g.toWord.length
  · simp only [startsWith, Set.mem_setOf_eq, getElem?_pos, Option.some.injEq,
      Prod.eq_iff_fst_eq_snd_eq, not_and, Bool.not_eq_not, toWord_mul, toWord_mk, reduce.cons,
      reduce_nil, List.cons_append, List.nil_append, reduce_toWord, hC] at *
    rw [show g.toWord = g.toWord.head (by grind) :: g.toWord.tail by grind]
    grind
  · simp_all [startsWith]

variable [MulAction (FreeGroup α) X]

instance {w : α × Bool} : SMul (startsWith w) X where
  smul g x := g.val • x

@[simp]
lemma startsWith.smul_def {w : α × Bool} {g : startsWith w} {x : X} : g • x = g.val • x := by
  rfl

/--
Applying `w⁻¹` to the orbit generated by all elements of a free group that start with `w` yields
the orbit generated by all the words that start with every letter except `w⁻¹`
(and the original point).
-/
theorem Orbit.duplicate (x : X) (w : α × Bool) :
    {(mk [w])⁻¹ • y | y ∈ MulAction.orbit (startsWith w) x} =
      (⋃ v ∈ {z : α × Bool | z ≠ (w.1, !w.2)}, MulAction.orbit (startsWith v) x) ∪ {x} := by
  ext i
  constructor
  · rintro ⟨-, ⟨⟨g, hg⟩, rfl⟩, rfl⟩
    set l := g.toWord with hl
    have h : (⟨g, hg⟩ : startsWith w) = ⟨mk g.toWord, by simp [g.mk_toWord, hg]⟩ := by
      simp [g.mk_toWord]
    match l with
    | [] => simp [← hl, startsWith] at hg
    | [a] =>
      simp_rw [h, ← hl, show a = w by simpa [← hl, startsWith] using hg, startsWith.smul_def,
        inv_smul_smul]
      exact Or.inr rfl
    | a :: b :: l =>
      have ha : a = w := by simpa [← hl, startsWith] using hg
      have h1 := isReduced_cons_cons.mp (hl ▸ isReduced_toWord)
      refine Or.inl (Set.mem_biUnion (x := b) (by grind) ?_)
      simp_rw [h, ← hl, ha, ← List.singleton_append (l := b :: l), ← mul_mk, startsWith.smul_def,
        mul_smul, inv_smul_smul]
      exact ⟨⟨mk (b :: l), by simp [startsWith, h1.2.reduce_eq]⟩, rfl⟩
  · rintro (⟨-, ⟨w', rfl⟩, -, ⟨hw, rfl⟩, ⟨g, hg⟩, rfl⟩ | rfl)
    · exact ⟨mk [w] • g • x, ⟨⟨mk [w] * g, startsWith_mk_mul g
        ((startsWith.disjoint_iff_ne.mpr hw).notMem_of_mem_left hg)⟩,
        mul_smul (mk [w]) g x⟩, inv_smul_smul (mk [w]) (g • x)⟩
    · exact ⟨mk [w] • i, ⟨⟨mk [w], rfl⟩, rfl⟩, inv_smul_smul (mk [w]) i⟩

end FreeGroup
